import torch.nn as nn
import numpy as np
from matplotlib import pyplot as plt
from data_loader import *
from cnn_train import *

# https://blog.csdn.net/ali1174/article/details/130294224
# https://blog.csdn.net/qq_39522016/article/details/130779099
# https://blog.csdn.net/weixin_43424450/article/details/129227788

class Adder(nn.Module):
    def __init__(self, cnn):
        super(Adder, self).__init__()
        self.cnn = cnn
        self.fc = nn.Sequential(
            nn.Linear(in_features=20, out_features=19),
            nn.Softmax()
        )
    def forward(self, x1, x2):
        out1 = self.cnn(x1)
        out2 = self.cnn(x2)
        out = torch.cat((out1, out2), dim=1)
        out = torch.relu(out)
        out = self.fc(out)
        return out


# 加载预训练模型
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
network=torch.load('./cnn.pth')

# 创建加法器模型
adder_model = Adder(network).to(device)

# 设置训练参数
epochs = 10

# 定义损失函数和优化器
criterion = nn.CrossEntropyLoss().to(device)
optimizer = torch.optim.AdamW(adder_model.parameters(), lr=0.003, weight_decay=0.002)

# 加载数据集
train_loader1,train_loader2,validation_loader1,validation_loader2=load_data_num_add()

# 训练加法器模型
all_train_loss=[]
val_acc=[]
for epoch in range(epochs):
    # 设置模型为训练模式
    adder_model.train()
    train_loss = 0
    val_loss=0
    for  (x_y1, x_y2) in zip(train_loader1, train_loader2):
        x1, y1 = [i.to(device) for i in x_y1]
        x2, y2 = [i.to(device) for i in x_y2]
        y = y1 + y2
        outputs = adder_model(x1, x2)
        loss = criterion(outputs, y)
        train_loss += loss.item()
        # 反向传播和优化
        optimizer.zero_grad()
        loss.backward()
        optimizer.step()

    all_train_loss.append(train_loss)
    # 打印训练信息
    print(f'迭代次数 {epoch+1}/{epochs}, Loss: {train_loss:.4f}',end='\t')

    # 进行验证
    total=0
    correct=0
    adder_model.eval()
    with torch.no_grad():  # 使用 no_grad() 的情况下，不计算梯度、不更新参数
        for  (va_x_y1, va_x_y2) in zip(validation_loader1, validation_loader2):
            x1, y1 = [i.to(device) for i in va_x_y1]
            x2, y2 = [i.to(device) for i in va_x_y2]
            y = y1 + y2
            outputs = adder_model(x1, x2)
            _, predicted = torch.max(outputs, dim=1)
            total += y.size(0)
            correct += (predicted == y).sum()
        validation_accuracy = 100 * correct / total
        print(f', 验证精度: {validation_accuracy:.2f}%')
        val_acc.append(validation_accuracy)





# # ----------------------数据可视化---------------------------------
#
# plt.figure(figsize=(10, 5))
# epochs_range = range(epochs)
# plt.plot(epochs_range, all_train_loss, label='Training Loss')
# plt.title('Training  Loss')
#
# plt.figure(figsize=(10, 5))
# plt.plot(epochs_range, val_acc, label='Validation Accuracy')
# plt.title('Validation Accuracy')
#
#
# plt.figure('数据可视化', figsize=(16, 5))
# for key,value in images_dict.items():
#     # 维度缩减
#     npimg = np.squeeze(value[0].numpy())
#     plt.subplot(2, 10, key)
#     plt.imshow(npimg, cmap=plt.cm.binary)
#     plt.axis('off')
# plt.show()

